import numpy as np
import pandas as pd
from scipy.optimize import minimize

def general_fit(X, y, model_func, loss_func, init_params, bounds=None, method='L-BFGS-B', options=None):
    """
    通用拟合接口
    X: 特征 (n_samples, n_features)
    y: 目标 (n_samples,)
    model_func: 用户自定义模型函数，格式为 f(X, params)
    loss_func: 用户自定义损失函数，格式为 loss(y_true, y_pred)
    init_params: 初始参数（数组）
    bounds: 参数边界
    method: 优化方法
    options: 优化器参数
    """
    def objective(params):
        y_pred = model_func(X, params)
        return loss_func(y, y_pred)
    
    res = minimize(objective, init_params, bounds=bounds, method=method, options=options)
    best_params = res.x
    y_pred = model_func(X, best_params)
    final_loss = loss_func(y, y_pred)
    return {
        'params': best_params,
        'y_pred': y_pred,
        'loss': final_loss,
        'success': res.success,
        'message': res.message
    } 

# 导入数据
data = pd.read_csv('data.csv', sep='\t')
# 新增dt列
data['dt'] = data['t_hf_wetbulb'] - data['outlet_water_temp']
# 生成 freq_all_flow 特征
data['freq_all_flow'] = data['freq_all'] / data['flow']


# 假设data已加载，并包含所需特征
x_cols = ['freq_all', 'flow', 'dt', 'freq_all_flow']
X = data[x_cols].values
y = data['heat_load'].values

# 线性模型函数：y_pred = w1*x1 + w2*x2 + ... + b
def model_func(X, params):
    c1, c2, epsilon1, epsilon2, epsilon3, d1,d2,d3,d4,d5,d6,d7,d8,d9,d10,d11,d12,n3 = params

    m_a = X[:, 0]  # freq_all
    m_chw = X[:, 1]  # flow
    t_s_ai = X[:, 2]  # t_hf_wetbulb
    t_wi = X[:, 3]    # outlet_water_temp
    numerator = c1 * (m_a ** epsilon1)+d1
    denominator = 1 + c2 * ((m_a / m_chw) ** epsilon2) +d2

    n1 = numerator / denominator * (t_s_ai - t_wi)**epsilon3 

    n2 = d1 * m_a**3 + d2 * (m_a / m_chw)**3 + d3 * (t_s_ai - t_wi)**3 \
        + d4 * m_a**2 * (m_a / m_chw) + d5 * (m_a / m_chw)**2* (t_s_ai - t_wi) + d6 * (t_s_ai - t_wi)**2 * m_a\
        + d7 * m_a* (m_a / m_chw)**2 + d8 * (m_a / m_chw)* (t_s_ai - t_wi)**2 + d9 * (t_s_ai - t_wi)* m_a**2 \
        + d10 * m_a** 2 + d11 * (m_a / m_chw)** 2 + d12 * (t_s_ai - t_wi)** 2
    
        
    return n1 +n3

# MSE损失函数
def loss_func(y_true, y_pred):
    return np.mean((y_true - y_pred) ** 2)

# 初始参数（特征数+1，最后一个为偏置）
init_params = np.zeros(18)

# 拟合
result = general_fit(X, y, model_func, loss_func, init_params)

print("最优参数：", result['params'])
print("最小损失：", result['loss'])
print("拟合是否成功：", result['success'])
print("优化信息：", result['message'])

# 增加评价指标
from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error, max_error
y_pred = result['y_pred']
r2 = r2_score(y, y_pred)
mae = mean_absolute_error(y, y_pred)
rmse = np.sqrt(mean_squared_error(y, y_pred))
mape = np.mean(np.abs((y - y_pred) / y)) * 100
maxerr = max_error(y, y_pred)

print(f"R²: {r2:.4f}")
print(f"MAE: {mae:.4f}")
print(f"RMSE: {rmse:.4f}")
print(f"MAPE: {mape:.2f}%")
print(f"Max Error: {maxerr:.4f}")